from helper import decode_json
from helper import *

import multiprocessing as mp
import threading
from queue import Queue as ThreadQueue
import time
import cv2
import os
import sys
import apriltag
import socket

sys.path.append("../")

from arm_control.angle import ArmAngle
from arm_control.move import ArmMove
from arm_control.func import execute_operation

CUBE_HEIGHT = 3.6
MAX_NUMBER_OF_OBEJECTS = 6

PICK_AND_RELEASE_OFFSET = 6
ARM_Z_OFFSET = 5

WINDOW_W = 640
WINDOW_H = 480
WINDOW_NAME = "Project P: Voice Control Robot Pick and Release"

SERVER_ADDRESS = "localhost"
SERVER_PORT = 12345
server_socket = None

CONFIG_FILE = os.path.expanduser("~/.bc_config/cam_calibration.config")
cfg = load_config(CONFIG_FILE)

arm_base_coordinate_x = cfg["arm_base_coordinate"]["x"]
arm_base_coordinate_y = cfg["arm_base_coordinate"]["y"]
pixel_per_cm_ratio = cfg["pixel_per_cm_ratio"]

# 放置区域范围(7，8,9在右边；10,11,12在左边)（x左正右负，y前正后负）单位cm）
place_position_in_pixels = {
    7: (arm_base_coordinate_x - int(18 * pixel_per_cm_ratio-50), arm_base_coordinate_y + int(4.5 * pixel_per_cm_ratio+10)),
    8: (arm_base_coordinate_x - int(18 * pixel_per_cm_ratio -45), arm_base_coordinate_y + int(17 * pixel_per_cm_ratio-20)),
    9: (arm_base_coordinate_x + int(18 * pixel_per_cm_ratio-30), arm_base_coordinate_y + int(4.5 * pixel_per_cm_ratio+10)),
    10: (arm_base_coordinate_x + int(18 * pixel_per_cm_ratio - 30), arm_base_coordinate_y + int(17 * pixel_per_cm_ratio-10)),
}

# 计算机械臂在摄像头像素坐标下距离原点的实际距离(x,y,z)
arm_base_pos = [
    *compute_pixel_to_dist(
        arm_base_coordinate_x, arm_base_coordinate_y, pixel_per_cm_ratio
    ),
    0,
]


def process_wrapper(queue, signal):
    def decorator(func):
        def wrapper():
            arm_move = ArmMove()
            arm_angle = ArmAngle(base_pos=arm_base_pos, z_offset=ARM_Z_OFFSET)
            while True:
                if not queue.empty():
                    targets = queue.get()
                    func(arm_move, arm_angle, *targets)
                    signal.set()
                time.sleep(0.01)  # Prevents busy waiting
        return wrapper
    return decorator


# Create a multiprocessing queue
queue = mp.Queue()
signal = mp.Event()
signal.set()

# 新增：用于缓存指令的队列
command_cache = ThreadQueue()

# Create a decorated version of pick_and_release
@process_wrapper(queue, signal)
def wrapped_execute_operation(*args):
    execute_operation(*args)


def socket_handler(sock, command_queue):
    # 处理socket接收到的指令，将其放入命令队列
    while True:
        data, addr = sock.recvfrom(1024)
        command = data.decode()
        command_queue.put(command)


def udp_server(host="localhost", port=12345):
    # 创建UDP服务器socket
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_socket.bind((host, port))
    return server_socket


def run(min_dist_threshold=10):
    # 主循环，处理摄像头画面、AprilTag检测、命令接收与机械臂控制
    process = mp.Process(target=wrapped_execute_operation)
    process.start()
    time.sleep(0.01)

    cap = cv2.VideoCapture(0)
    detector = apriltag.Detector()

    cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(WINDOW_NAME, WINDOW_W, WINDOW_H)

    place_position_in_dist = {
        name: (*compute_pixel_to_dist(x, y, pixel_per_cm_ratio), CUBE_HEIGHT)
        for name, (x, y) in place_position_in_pixels.items()
    }

    print(place_position_in_dist)

    target_box_size_in_pixel = int(CUBE_HEIGHT * pixel_per_cm_ratio)

    server_socket = udp_server(host=SERVER_ADDRESS, port=SERVER_PORT)

    command_queue = ThreadQueue()

    threading.Thread(
        target=socket_handler, args=(server_socket, command_queue), daemon=True
    ).start()

    i = 0
    cached_detected_pos = {}
    cached_goal_pos = -1
    pos_dic = {}

    try:
        ### 整个while 是检测函数的主体, 在主体里面插入 pick_and_release 函数进程的控制信号量
        while True:
            ret, frame = cap.read()
            gray_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            # 对输入的图像检测 Apriltag
            tag_pos_pairs = []
            results = detector.detect(gray_img)
            for result in results:
                tag_id = result.tag_id
                top_left, top_right, bottom_right, bottom_left = result.corners
                center_x = int((top_left[0] + bottom_right[0]) / 2)
                center_y = int((top_left[1] + bottom_right[1]) / 2)
                tag_inform = (tag_id, center_x, center_y)
                tag_pos_pairs.append(tag_inform)

            # 可视化检测结果
            for tag in tag_pos_pairs:
                draw_tag(frame, tag[0], tag[1:])
                cached_detected_pos[tag[0]] = (*compute_pixel_to_dist(tag[1], tag[2], pixel_per_cm_ratio), CUBE_HEIGHT)

            # 可视化放置位置
            frame = draw_target_pos(frame, place_position_in_pixels, width=target_box_size_in_pixel)
            cv2.putText(
                frame,
                f"Jobs Count: {i}",
                (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX,
                1,
                (0, 255, 0),
                2,
            )

            # 从信息队列中读取控制指令
            while not command_queue.empty():
                command = command_queue.get()
                print(f"\n接收到的command:\n {command}")
                # 如果机械臂正在执行任务，将指令缓存起来
                if not signal.is_set():
                    command_cache.put(command)
                    print(f"机械臂正在任务中, 指令已缓存，待任务完成后处理。")
                else:
                    command = decode_json(command)
                    goal_pos = command["target"]  # 目标位置解码
                    operation = command["intent"]  # 目标动作解码
                    # 目标位置在 1. 算法的检测结果缓存队列中   2. 程序预先定义的固定位置   3. 没有给目标位置
                    if goal_pos in cached_detected_pos or goal_pos in place_position_in_dist or goal_pos == -1:
                        print(
                            f"将坐标发送到队列: \n\t当前目标为 {goal_pos}\n\t操作类型为 {operation}"
                        )
                        if goal_pos in cached_detected_pos:
                            pos_dic = cached_detected_pos
                            cached_goal_pos = goal_pos
                        elif goal_pos in place_position_in_dist:
                            pos_dic = place_position_in_dist
                            cached_goal_pos = goal_pos
                        else:
                            goal_pos = cached_goal_pos
                            print(f"使用缓存的{goal_pos}位置")
                        # 当用户说 " 到xx位置去",  后面跟随的很有可能是 "放置", "抓取" 等没有目标位置。所以要缓存'xx'
                        if operation == "go_to":
                            cached_goal_pos = goal_pos
                        if goal_pos == -1:
                            continue
                        queue.put(
                            (
                                pos_dic[goal_pos][0],
                                pos_dic[goal_pos][1],
                                pos_dic[goal_pos][2],
                                operation,
                                PICK_AND_RELEASE_OFFSET,
                            )
                        )
                    else:
                        print(f"没有找到任务目标: {goal_pos}")
                        continue
                    i += 1
                    signal.clear()

            # 机械臂任务完成后，处理缓存的指令
            if signal.is_set() and not command_cache.empty():
                command = command_cache.get()
                print(f"\n处理缓存的command:\n {command}")
                command = decode_json(command)
                goal_pos = command["target"]  # 目标位置解码
                operation = command["intent"]  # 目标动作解码
                # 目标位置在 1. 算法的检测结果缓存队列中   2. 程序预先定义的固定位置   3. 没有给目标位置
                if goal_pos in cached_detected_pos or goal_pos in place_position_in_dist or goal_pos == -1:
                    print(
                        f"将坐标发送到队列: \n\t当前目标为 {goal_pos}\n\t操作类型为 {operation}"
                    )
                    if goal_pos in cached_detected_pos:
                        pos_dic = cached_detected_pos
                        cached_goal_pos = goal_pos
                    elif goal_pos in place_position_in_dist:
                        pos_dic = place_position_in_dist
                        cached_goal_pos = goal_pos
                    else:
                        goal_pos = cached_goal_pos
                        print(f"使用缓存的{goal_pos}位置")
                    # 当用户说 " 到xx位置去",  后面跟随的很有可能是 "放置", "抓取" 等没有目标位置。所以要缓存'xx'
                    if operation == "go_to":
                        cached_goal_pos = goal_pos
                    if goal_pos == -1:
                        continue
                    queue.put(
                        (
                            pos_dic[goal_pos][0],
                            pos_dic[goal_pos][1],
                            pos_dic[goal_pos][2],
                            operation,
                            PICK_AND_RELEASE_OFFSET,
                        )
                    )
                else:
                    print(f"没有找到任务目标: {goal_pos}")
                    continue
                i += 1
                signal.clear()

            cv2.imshow(WINDOW_NAME, frame)
            if cv2.waitKey(10) & 0xFF == ord('q'):
                break
    except KeyboardInterrupt:
        print("Process interrupted by user")
    finally:
        process.terminate()
        process.join()


if __name__ == "__main__":
    min_dist_threshold = 15
    run(min_dist_threshold)
